#pragma once

#include <array>
#include <cstddef>

#include "../utils/misc.h"

namespace mdtx::online::impl
{
    template <typename T>
    class sum
    {
    public:
        T s;
        kahan<T> acc_s;

        void insert_impl(T x)
        {
            acc_s(s, x);
        }

        void remove_impl(T old_x)
        {
            acc_s(s, -old_x);
        }

        void roll_impl(T x, T old_x)
        {
            acc_s(s, x - old_x);
        }

        void reset_impl()
        {
            s = 0;
            acc_s.reset();
        }
    };

    template <typename T>
    class sma
    {
    public:
        ptrdiff_t n;
        T m;
        kahan<T> acc_m;

        void insert_impl(T x)
        {
            T d = x - m;
            n += 1;
            acc_m(m, d / n);
        }

        void remove_impl(T old_x)
        {
            T d = m - old_x;
            n -= 1;
            acc_m(m, d / n);
        }

        void roll_impl(T x, T old_x)
        {
            T d = x - old_x;
            acc_m(m, d / n);
        }

        void reset_impl()
        {
            n = 0;
            m = 0;
            acc_m.reset();
        }
    };

    template <typename T>
    class wma
    {
    public:
        ptrdiff_t n, tot_w;
        T xsum, wxsum;
        kahan<T> acc_x, acc_wx;

        void insert_impl(T x)
        {
            n += 1;
            tot_w += n;
            acc_wx(wxsum, x * n);
            acc_x(xsum, x);
        }

        void remove_impl(T old_x)
        {
            tot_w -= n;
            n -= 1;
            acc_wx(wxsum, -xsum);
            acc_x(xsum, -old_x);
        }

        void roll_impl(T x, T old_x)
        {
            acc_wx(wxsum, x * n - xsum);
            acc_x(xsum, x - old_x);
        }

        void reset_impl()
        {
            n = tot_w = 0;
            xsum = wxsum = 0;
            acc_x.reset();
            acc_wx.reset();
        }
    };

    template <typename T>
    class wsma
    {
    public:
        T m, wsum;
        kahan<T> acc_m, acc_w;

        void insert_impl(T x, T w)
        {
            T d = x - m;
            acc_w(wsum, w);
            acc_m(m, (w / wsum) * d);
        }

        void remove_impl(T old_x, T old_w)
        {
            T d = m - old_x;
            acc_w(wsum, -old_w);
            acc_m(m, (old_w / wsum) * d);
        }

        void reset_impl()
        {
            m = wsum = 0;
            acc_m.reset();
            acc_w.reset();
        }
    };

    template <typename T>
    class ewma
    {
    public:
        T a, a1, a1n;
        T wsum{}, tot_w{};
        kahan<T> acc_tot_w{};

        /**
         * @brief Construct a new ewma impl object
         *
         * @param a exponential weight factor
         */
        explicit ewma(T alpha)
            : a(alpha), a1(1 - alpha), a1n(1)
        {
        }

        void insert_impl(T x)
        {
            acc_tot_w(tot_w, a1n);
            wsum = a1 * wsum + x;
            a1n *= a1;
        }

        void remove_impl(T old_x)
        {
            a1n /= a1;
            wsum -= a1n * old_x;
            acc_tot_w(tot_w, -a1n);
        }

        void roll_impl(T x, T old_x)
        {
            wsum = a1 * wsum + x - a1n * old_x;
        }

        void reset_impl()
        {
            a1 = 1 - a;
            a1n = 1;
            wsum = tot_w = 0;
            acc_tot_w.reset();
        }
    };

    template <typename T>
    class ema
    {
    public:
        T a, m{};
        kahan<T> acc_m{};
        bool init{};

        /**
         * @brief Construct a new ema impl object
         *
         * @param alpha exponential weight factor
         */
        explicit ema(T alpha)
            : a(alpha)
        {
        }

        void insert_impl(T x)
        {
            if (init)
                [[likely]]
            {
                acc_m(m, (x - m) * a);
            }
            else
                [[unlikely]]
            {
                m = x;
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            m = 0;
            acc_m.reset();
        }
    };

    template <typename T>
    class ema_evw
    {
    public:
        T m;
        kahan<T> acc_m;
        bool init;

        void insert_impl(T x, T w)
        {
            if (init)
                [[likely]]
            {
                acc_m(m, (x - m) * w);
            }
            else
                [[unlikely]]
            {
                m = x;
                init = true;
            }
        }

        void reset_impl()
        {
            m = 0;
            acc_m.reset();
        }
    };

    template <typename T, size_t N>
    class xema
    {
    public:
        T a;
        // mean
        std::array<T, N> m{};
        std::array<kahan<T>, N> acc_m{};
        bool init{};

        /**
         * @brief Construct a new xema updater object
         *
         * @param alpha exponential weight factor
         */
        explicit xema(T alpha)
            : a(alpha)
        {
        }

        void insert_impl(T x)
        {
            if (init)
                [[likely]]
            {
                acc_m[0](m[0], (x - m[0]) * a);
                for (size_t i = 1; i < N; ++i)
                {
                    acc_m[i](m[i], (m[i - 1] - m[i]) * a);
                }
            }
            else
                [[unlikely]]
            {
                m.fill(x);
                init = true;
            }
        }

        void reset_impl()
        {
            init = false;
            m.fill(0);
            for (auto &acc : acc_m)
                acc.reset();
        }
    };

    template <typename T, size_t N>
    class xsma
    {
    public:
        ptrdiff_t n;
        std::array<T, N> m;
        std::array<kahan<T>, N> acc_m;

        template <typename... Args>
        void insert_impl(Args... x)
        {
            n += 1;
            size_t i{}, j{}, k{};
            ((acc_m[i++](m[j++], (x - m[k++]) / n)), ...);
        }

        template <typename... Args>
        void remove_impl(Args... x)
        {
            n -= 1;
            size_t i{}, j{}, k{};
            ((acc_m[i++](m[j++], (m[k++] - x) / n)), ...);
        }

        template <typename... Args>
        void roll_impl(Args... x)
        {
            const std::array<T, N * 2> xs{x...};
            for (size_t i = 0; i < N; ++i)
            {
                acc_m[i](m[i], (xs[i] - xs[i + N]) / n);
            }
        }

        void reset_impl()
        {
            m.fill(0);
            for (auto &acc : acc_m)
                acc.reset();
        }
    };
}
